import numpy as np
import matplotlib.pyplot as plt
import math
from osgeo import gdal
gdal.AllRegister()


def read_img(filename):
    dataset = gdal.Open(filename)
    ncols = dataset.RasterXSize  # 栅格矩阵的列数
    nrows = dataset.RasterYSize  # 栅格矩阵的行数
    proj = dataset.GetProjection()  # 地图投影信息，字符串表示
    data = dataset.ReadAsArray(0, 0)
    datatype = data.dtype
    adfGeoTransform = dataset.GetGeoTransform()
    # DEM的平面四至
    Xmin = adfGeoTransform[0]
    Ymin = adfGeoTransform[3]
    Xmax = adfGeoTransform[0] + nrows * adfGeoTransform[1] \
           + ncols * adfGeoTransform[2]
    Ymax = adfGeoTransform[3] + nrows * adfGeoTransform[4] \
           + ncols * adfGeoTransform[5]
    xcellwidth = abs(adfGeoTransform[1])
    ycellwidth = abs(adfGeoTransform[5])
    del dataset  # 关闭对象dataset，释放内存
    return {"nrows": nrows, "ncols": ncols, "data": data,
            "Xmin": Xmin, "Ymin": Ymin, "Xmax": Xmax, "Ymax": Ymax,
            "xcellwidth": xcellwidth, "ycellwidth": ycellwidth,
            "proj": proj, "datatype": datatype}


def CalculateSlope(demMatrix, slopeMatrix, x, y, xcellsize=30, ycellsize=30, neighbors=3):
    """
    计算坡度
    Args:
        demMatrix (ndarray): dem矩阵外圈拓展后的矩阵
        slopeMatrix (ndarray): 生成的slope矩阵
        x (int): x坐标
        y (int): y坐标
        xcellsize (int, optional): [description]. Defaults to 30.
        ycellsize (int, optional): [description]. Defaults to 30.
        neighbors (int, optional): [description]. Defaults to 3.
    """
    searchRadius = neighbors // 2
    # 坡度窗口
    slopeWindow = demMatrix[x:x + 2 * searchRadius + 1, y:y + 2 * searchRadius + 1]
    _ySlopeRate1 = (slopeWindow[2, 0] + slopeWindow[2, 2] + slopeWindow[2, 1] * 2)
    _ySlopeRate2 = (slopeWindow[0, 0] + slopeWindow[0, 2] + slopeWindow[0, 1] * 2)
    ySlopeRate = (_ySlopeRate2 - _ySlopeRate1) / 8 / xcellsize
    xSlopeRate = ((slopeWindow[2, 2] + slopeWindow[0, 2] + slopeWindow[1, 2] * 2)
                  -
                  (slopeWindow[2, 0] + slopeWindow[0, 0] + slopeWindow[1, 0] * 2)) / 8 / ycellsize
    rise_run = math.sqrt(xSlopeRate * xSlopeRate + ySlopeRate * ySlopeRate)
    slope_degrees = math.atan(rise_run) * 180 / math.pi
    slopeMatrix[x, y] = slope_degrees



def Execution(basepath, filename, extended=True, visible=True):
    filepath = basepath + "/" + filename
    imgData = read_img(filepath)
    dem = imgData["data"]
    if extended:
        dem = np.insert(dem, 0, dem[0], axis=0)
        dem = np.vstack([dem, dem[imgData['nrows']]])
        dem = np.insert(dem, 0, dem[:, 0], axis=1)
        dem = np.insert(dem, imgData['ncols'],
                        dem[:, imgData['ncols']], axis=1)
    else:
        pass
    SlopeMatrix = np.empty([imgData["nrows"], imgData["ncols"]], dtype=float)
    for i in range(0, imgData["nrows"]):
        for j in range(0, imgData["ncols"]):
            CalculateSlope(dem, SlopeMatrix, i, j)
    if visible:
        x = np.linspace(imgData['Xmin'] - imgData['xcellwidth'],
                        imgData['Xmax'] + imgData['xcellwidth'], imgData["ncols"])
        y = np.linspace(imgData['Ymin'] - imgData['ycellwidth'],
                        imgData['Ymax'] + imgData['ycellwidth'], imgData["nrows"])
        X, Y = np.meshgrid(x, y)
        region = np.s_[10:400, 10:400]
        X, Y = X[region], Y[region]
        slope = SlopeMatrix[10:400, 10:400]
        plt.pcolormesh(X, Y, slope)
        plt.colorbar()
        plt.savefig(filepath + "/"+"pic.png")  # 保存图片
    else:
        return SlopeMatrix


# 按间距中的绿色按钮以运行脚本。
if __name__ == '__main__':
    Execution("/demcode-learning/dem", 'Himalaya.tif')
